import matplotlib.pyplot as plt
import numpy as np
import util

from linear_model import LinearModel


def main(tau, train_path, eval_path):
    """Problem 5(b): Locally weighted regression (LWR)

    Args:
        tau: Bandwidth parameter for LWR.
        train_path: Path to CSV file containing dataset for training.
        eval_path: Path to CSV file containing dataset for evaluation.
    """
    # Load training set
    x_train, y_train = util.load_dataset(train_path, add_intercept=True)

    # *** START CODE HERE ***
    # Fit a LWR model
    clf = LocallyWeightedLinearRegression(tau)
    clf.fit(x_train, y_train)

    # Load validation set
    x_eval, y_eval = util.load_dataset(eval_path, add_intercept=True)

    # Make predictions
    y_pred = clf.predict(x_eval)

    # Get MSE value on the validation set
    mse = np.mean((y_pred - y_eval) ** 2)
    print(f"Locally Weighted Regression MSE: {mse:.4f}")

    # Plot validation predictions on top of training set
    plt.figure()
    plt.plot(x_train[:, 1], y_train, 'bx', label='Training Data')
    plt.plot(x_eval[:, 1], y_pred, 'ro', label='Validation Predictions')
    plt.xlabel('x')
    plt.ylabel('y')
    plt.title(f'Locally Weighted Regression (τ={tau})')
    plt.legend()
    plt.savefig('output/p05b.png')
    plt.show()
    # *** END CODE HERE ***


class LocallyWeightedLinearRegression(LinearModel):
    """Locally Weighted Regression (LWR).

    Example usage:
        > clf = LocallyWeightedLinearRegression(tau)
        > clf.fit(x_train, y_train)
        > clf.predict(x_eval)
    """

    def __init__(self, tau):
        super(LocallyWeightedLinearRegression, self).__init__()
        self.tau = tau
        self.x = None
        self.y = None

    def fit(self, x, y):
        """Fit LWR by saving the training set.

        """
        # *** START CODE HERE ***
        self.x = x
        self.y = y
        # *** END CODE HERE ***

    def predict(self, x):
        """Make predictions given inputs x.

        Args:
            x: Inputs of shape (m, n).

        Returns:
            Outputs of shape (m,).
        """
        # *** START CODE HERE ***
        m = x.shape[0]
        predictions = np.zeros(m)

        for i in range(m):

            distances = np.sum((self.x[:, 1:] - x[i, 1:]) ** 2, axis=1)
            weights = np.exp(-distances / (2 * self.tau ** 2))

            # Create diagonal weight matrix W
            W = np.diag(weights)
            X = self.x
            y = self.y

            XTWX = X.T @ W @ X
            XTWy = X.T @ W @ y

            theta = np.linalg.solve(XTWX, XTWy)
            predictions[i] = x[i] @ theta

        return predictions
        # *** END CODE HERE ***
